/**
* Copyright 2020 Huawei Technologies Co., Ltd
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at

* http://www.apache.org/licenses/LICENSE-2.0

* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.

* File sample_process.cpp
* Description: handle acl resource
*/
#include "deepsort/deep_sort.h"
#include <iostream>
#include "acl/acl.h"
#include "utils.h"
#include <opencv2/videoio.hpp>

using namespace std;


DeepSort::DeepSort(const char* modelPath,
uint32_t modelWidth,
uint32_t modelHeight)
    :extractor(make_shared<Extractor>(modelPath, modelWidth, modelHeight)),
    manager(make_shared<TrackerManager<TrackData>>(data)),
    feat_metric(make_shared<FeatureMetric<TrackData>>(data)),
    drawer(make_shared<MotDrawer>()),
    isInited_(false) {}

DeepSort::~DeepSort() {
    DestroyResource();
}


Result DeepSort::Init(int32_t& deviceId, aclrtContext& context, aclrtStream& stream, aclrtRunMode& runMode) {
    if (isInited_) {
        INFO_LOG("DeepSort instance is initied already!");
        return SUCCESS;
    }
    extractor->Init(deviceId, context, stream, runMode);


    isInited_ = true;
    return SUCCESS;
}

Result DeepSort::Process(cv::Mat& ori_img, 
            std::vector<BBox>& detections, std::vector<BBox>& trackResults) {
    
    manager->predict();
    manager->remove_nan();

    vector<vector<float>> feature;
    auto matched = manager->update(
        detections,
        [this, &detections, &ori_img, &feature](const std::vector<int> &trk_ids, const std::vector<int> &det_ids) -> vector<vector<float>> {
            vector<BBox> trks;
            for (auto t : trk_ids) {
                BBox b;
                b.rect = data[t].kalman.rect();
                trks.push_back(b);
            }
            vector<BBox> dets;
            for (auto d:det_ids) {
                dets.push_back(detections[d]);
            }

            auto iou_mat = iou_dist(dets, trks);

            // 15ms
            Result ret = extractor->extract(ori_img, dets, feature);
            if (ret != SUCCESS) {
                ERROR_LOG("extract feature fail!");
            }
            auto feat_mat = feat_metric->distance(feature, trk_ids);
            for (size_t i = 0; i < feat_mat.size(); ++i) {
                for (size_t j = 0; j < feat_mat[0].size(); ++j) {
                    if (iou_mat[i][j]>0.7f || feat_mat[i][j]>0.05f) {
                        feat_mat[i][j] = INVALID_DIST;
                    }
                    if (feat_mat[i][j] < 0) {
                        feat_mat[i][j] = 0;
                    }
                }
            }

            return feat_mat;
        },
        [this, &detections](const std::vector<int> &trk_ids, const std::vector<int> &det_ids) -> vector<vector<float>> {
            vector<BBox> trks;
            for (auto t : trk_ids) {
                BBox b;
                b.rect = data[t].kalman.rect();
                trks.push_back(b);
            }
            vector<BBox> dets;
            for (auto &d:det_ids) {
                dets.push_back(detections[d]);
            }
            auto iou_mat = iou_dist(dets, trks);
            for (size_t i = 0; i < iou_mat.size(); ++i) {
                for (size_t j = 0; j < iou_mat[0].size(); ++j) {
                    if (iou_mat[i][j] > 0.6f) {
                        iou_mat[i][j] = INVALID_DIST;
                    }
                }
            }
            return iou_mat;
        });

    vector<int> targets;
    vector<vector<float>> matched_feature;
    for (auto[x, y]:matched) {
        targets.emplace_back(x);
        matched_feature.emplace_back(feature[y]);
    }

    feat_metric->update(matched_feature, targets);
    manager->remove_deleted();
    manager->visible_tracks(trackResults);



    // Utils::WriteImageToFile(ori_img, "output_deepsort");

    return SUCCESS;
}

void DeepSort::DrawTrajectories(cv::Mat& ori_img, vector<BBox>& trackResults, bool isDrawTrajectory) 
{
    drawer->DrawTrajectories(ori_img, trackResults, isDrawTrajectory);
}



Result DeepSort::Set(uint32_t width, uint32_t height) {
    manager->Set(width, height);
    return SUCCESS;
}

void DeepSort::DestroyResource() {
    extractor->DestroyResource();
}